#include <stdlib.h>
#include <string.h>
#include "bignum.h"

uint8_t hex_to_byte(const char ch) {
    if (ch >= '0' && ch <= '9') {
        return ch - '0';
    } else if (ch >= 'A' && ch <= 'F') {
        return ch - 'A' + 10;
    } else if (ch >= 'a' && ch <= 'f') {
        return ch - 'a' + 10;
    } else {
        return 16;
    }
}

int bignum_set_hex(bignum_t* num, const char* str) {
    int i = strlen(str), n = 0, size;
    if (i == 0) {
        return 1;
    }
    size = i / 2 + i % 2;
    uint8_t* buffer = malloc(size);
    if (buffer == NULL) {
        return 1;
    }
    num->data = buffer;
    while (n < size) {
        i -= 2;
        uint8_t byte, high, low;
        if (i >= 0) {
            high = hex_to_byte(str[i]);
        } else {
            high = 0;
        }
        low = hex_to_byte(str[i + 1]);
        if (high > 15 || low > 15) {
            free(buffer);
            return 1;
        }
        byte = (high << 4) | low;
        buffer[n++] = byte;
    }
    num->size = n;
    return 0;
}

int bignum_set_str(bignum_t* num, const char* str) {
    int i = strlen(str);
    num->size = i;
    num->data = malloc(i);
    if (num->data == NULL) {
        return 1;
    }
    while (i > 0) {
        num->data[num->size - i] = str[i - 1];
        i--;
    }
    return 0;
}

char num_to_hex(uint8_t num) {
    if (num >= 16) {
        return 0;
    }
    if (num < 10) {
        return '0' + num;
    } else {
        return 'a' + num - 10;
    }
}

char* bignum_to_hex(bignum_t* bignum){
    char* result = malloc(bignum->size * 2 + 1);
    if (result == NULL) {
        return NULL;
    }

    for (int i = 0; i < bignum->size; i++) {
        uint8_t byte = bignum->data[bignum->size - 1 - i];
        uint8_t high = byte >> 4;
        uint8_t low = byte & 0xF;
        char high_char = num_to_hex(high);
        char low_char = num_to_hex(low);
        if (high_char == 0 || low_char == 0) {
            free(result);
            return NULL;
        }
        result[i * 2] = high_char;
        result[i * 2 + 1] = low_char;
    }

    result[bignum->size * 2] = '\0';
    if (result[0] == '0') {
        char* trimmed = malloc(bignum->size * 2);
        strcpy(trimmed, result + 1);
        free(result);
        return trimmed;
    }
    return result;
}

char* encrypt_password(const char* password, const char* exponent, const char* modulus) {
    bignum_t password_bignum, modulus_bignum, result_bignum;
    if (bignum_set_str(&password_bignum, password)) {
        return NULL;
    }
    if (bignum_set_hex(&modulus_bignum, modulus)) {
        return NULL;
    }

    char* end = (char*) exponent;
    uint32_t exponent_num = strtol(exponent, &end, 16);
    if (*end != '\0') {
        return NULL;
    }

    bignum_pow_mod(&password_bignum, exponent_num, &modulus_bignum, &result_bignum);

    bignum_free(&password_bignum);
    bignum_free(&modulus_bignum);
    bignum_trim(&result_bignum);
    return bignum_to_hex(&result_bignum);
}
